When the sampling strategy is "shuffle", does this shuffle the entire dataset before indexing, or does it shuffle the indexed array?

I think it shuffles the indexed array because the where clause is most likely applied to filter first.

Is that what we want? Seems like "shuffle" should mean "shuffle the entire dataset" so that all records are equally likely to be present in a given indexed array. I'm not sure it actually matters too much, though, since all records are treated independently – assuming the entire seed dataset is used.

https://duckdb.org/docs/stable/sql/query_syntax/orderby

It's applied after the where clause filtering.

It would be weird to shuffle first and then give exact indices ....

Let's think about the user's desired outcomes:

• Shuffled, no repeat: "Just give me one entry per row, I don't care about anything else"
• Shuffled, repeat: "I just want a bunch of random outcomes, I'm using my inputs as another sampler."
• Ordered, no repeat: "I NEED a one-to-one, in order correspondence with this data."
• Ordered, repeat: "I want multiple epochs/passes over this data, and the same number of them per sample."

How this relates to PartitionBlock:

• Shuffled, no repeat: Split into partitions, randomize order within partition, no problem.
• Shuffled, repeat: Depends on how exactly we want to preserve statistics. True stat would be shuffle and then split. If the partitions were exactly the same size, then it would be okay, but that is not true in general (esp. for the last Partition which might have some remainder size). However, there is a fix -- we can change the number of samples produced from each PartitionBlock according to the proportion of the size of the Block with respect to the full dataset. That should converge us to the same distribution of the original shuffle.
• Ordered, no repeat: Split into partitions, create a record for each row, no problem.
• Ordered, repeat: Split into partitions, create N records for each row, no problem.

I think that doing the proportional sampling per PartitionBlock post-partitioning would be desirable since it parallelizes easily vs needing to do an op on the full dataset.

The exact partition that is chosen is hidden from the user.

We resolve to indices automatically for the user, but it's in order based on how they define partitions.

For posterity, we chatted on Slack about this. To solve for a weird distribution getting stuck in the the last partition ( the case of shuffle), we'd need to pre-shuffle the whole dataset. It is an expensive operation.

I ran a test to see how long it would take for duckdb to shuffle a really large dataset. As an example, we took the web step dataset that's about 164 GB.

TL;DR
duckdb took about a half hour to shuffle the entire dataset and write to 1000 partitions in a different destination:

srun Execution

srun --account=llmservice_sdg_research --partition=cpu --nodes=1 --ntasks=1 --cpus-per-task=8 --mem=128G --time=2:00:00 bash -c "source /lustre/fsw/portfolios/llmservice/users/nmulepati/test_shuffle/.venv/bin/activate && python shuffle.py"

Output

Starting shuffle and write...
Configuring DuckDB settings...
Reading and shuffling data...
Source: /lustre/fsw/portfolios/llmservice/users/etramel/datasets/stem_pretrain/dataprep/v2_16k_datamix/*.parquet
Destination: /lustre/fsw/portfolios/llmservice/users/nmulepati/test_shuffle/tmp_data
Initial partitioning complete.
Shuffle and write completed in 1733.94 seconds
Output written to: /lustre/fsw/portfolios/llmservice/users/nmulepati/test_shuffle/tmp_data

shuffle.py

import duckdb
import time

start = time.perf_counter()
print("Starting shuffle and write...", flush=True)

# Connect to DuckDB
conn = duckdb.connect()

# Configure DuckDB for memory-efficient operation
# Conservative memory settings to avoid OOM
print("Configuring DuckDB settings...", flush=True)
conn.execute("SET memory_limit='110GB'")  # Leave headroom for system
conn.execute("SET threads=4")  # Reduce threads to save memory
conn.execute("SET temp_directory='/tmp/duckdb_temp'")
conn.execute("SET preserve_insertion_order=false")
# Enable disk-based operations for large datasets
conn.execute("SET max_memory='110GB'")

# Source and destination paths
data_path = "/lustre/fsw/portfolios/llmservice/users/etramel/datasets/stem_pretrain/dataprep/v2_16k_datamix/*.parquet"
output_path = "/lustre/fsw/portfolios/llmservice/users/nmulepati/test_shuffle/tmp_data"

print("Reading and shuffling data...", flush=True)
print(f"Source: {data_path}", flush=True)
print(f"Destination: {output_path}", flush=True)

# Memory-efficient shuffle: just add random partition key, no ORDER BY
# The random partitioning itself provides the shuffle effect
query = f"""
    COPY (
        SELECT *, (random() * 1000)::INTEGER as _partition_key
        FROM '{data_path}'
    ) TO '{output_path}' (
        FORMAT PARQUET,
        PARTITION_BY (_partition_key),
        OVERWRITE_OR_IGNORE true
    )
"""

conn.execute(query)
print("Initial partitioning complete.", flush=True)

elapsed = time.perf_counter() - start
print(f"Shuffle and write completed in {elapsed:.2f} seconds")
print(f"Output written to: {output_path}")

Ok I've optimized the select query to be based on limit and offset followed by order by. The previous implementation was still doing a full table scan. 483363d

Example run to pull a 1000 item slice from the 164GB partitioned dataset took <1s:

Starting select...
Start: 435971.227842381
End: 10999
Limit: 1000
Offset: 10000
Query:

SELECT * from (
    SELECT * from '/lustre/fsw/portfolios/llmservice/users/etramel/datasets/stem_pretrain/dataprep/v2_16k_datamix/*.parquet' LIMIT 1000 OFFSET 10000
) ORDER BY RANDOM()

Executing query...
Fetched 1000 rows
Select completed in 0.88 seconds

cc @eric-tramel @johnnygreco